import asyncio  # used to run async functions
# Patch openai to use our mock
import sys
import types
from typing import List

import pytest  # used for our unit tests
from guardrails.utils.openai_utils.base import BaseOpenAIClient
from guardrails.utils.openai_utils.v1 import AsyncOpenAIClientV1

# --- Begin: Minimal stubs for base client ---

class BaseOpenAIClient:
    def __init__(
        self,
        api_key: str = None,
        api_base: str = None,
    ):
        self.api_key = api_key
        self.api_base = api_base
from guardrails.utils.openai_utils.v1 import AsyncOpenAIClientV1

# --- Begin: Unit Tests ---

@pytest.mark.asyncio

















#------------------------------------------------
import asyncio  # used to run async functions
from typing import List

import openai
import pytest  # used for our unit tests
from guardrails.utils.openai_utils.base import BaseOpenAIClient
from guardrails.utils.openai_utils.v1 import AsyncOpenAIClientV1

# --- Unit Tests ---

# We'll use pytest's monkeypatch to mock openai.AsyncClient for deterministic unit tests.
# This avoids external calls and lets us test async behavior, edge cases, and throughput.

class DummyEmbedding:
    def __init__(self, embedding):
        self.embedding = embedding

class DummyEmbeddingsResponse:
    def __init__(self, embeddings):
        # embeddings is a list of lists of floats
        self.data = [DummyEmbedding(e) for e in embeddings]

class DummyAsyncEmbeddingsCreate:
    def __init__(self, embeddings_map):
        self.embeddings_map = embeddings_map  # Dict of (model, input) -> embeddings

    async def create(self, model, input):
        # Simulate async embedding creation
        # If input is empty, return empty list
        if not input:
            return DummyEmbeddingsResponse([])
        # If model or input not in map, return dummy values
        key = (model, tuple(input))
        if key in self.embeddings_map:
            return DummyEmbeddingsResponse(self.embeddings_map[key])
        # Default: return a vector of [0.0, ...] for each input string
        return DummyEmbeddingsResponse([[0.0] * 3 for _ in input])

class DummyAsyncClient:
    def __init__(self, embeddings_map=None, api_key=None, base_url=None):
        self.embeddings = DummyAsyncEmbeddingsCreate(embeddings_map or {})

# Helper to patch openai.AsyncClient with our dummy
@pytest.fixture
def patch_openai_async_client(monkeypatch):
    def _patch(embeddings_map=None):
        monkeypatch.setattr(openai, "AsyncClient", lambda api_key=None, base_url=None: DummyAsyncClient(embeddings_map, api_key, base_url))
    return _patch

# ---- BASIC TEST CASES ----

@pytest.mark.asyncio
async def test_create_embedding_basic_single(monkeypatch, patch_openai_async_client):
    """Test basic async/await behavior with a single input string."""
    embeddings_map = {
        ("test-model", ("hello world",)): [[0.1, 0.2, 0.3]]
    }
    patch_openai_async_client(embeddings_map)
    client = AsyncOpenAIClientV1(api_key="dummy-key")
    result = await client.create_embedding("test-model", ["hello world"])

@pytest.mark.asyncio
async def test_create_embedding_basic_multiple(monkeypatch, patch_openai_async_client):
    """Test basic async/await behavior with multiple input strings."""
    embeddings_map = {
        ("test-model", ("foo", "bar")): [[0.5, 0.6, 0.7], [0.8, 0.9, 1.0]]
    }
    patch_openai_async_client(embeddings_map)
    client = AsyncOpenAIClientV1(api_key="dummy-key")
    result = await client.create_embedding("test-model", ["foo", "bar"])

@pytest.mark.asyncio

async def test_create_embedding_edge_long_strings(monkeypatch, patch_openai_async_client):
    """Test with very long input strings."""
    long_str = "a" * 10000
    embeddings_map = {
        ("model-x", (long_str,)): [[0.1, 0.2, 0.3]]
    }
    patch_openai_async_client(embeddings_map)
    client = AsyncOpenAIClientV1(api_key="dummy-key")
    result = await client.create_embedding("model-x", [long_str])

@pytest.mark.asyncio

async def test_create_embedding_edge_concurrent(monkeypatch, patch_openai_async_client):
    """Test concurrent execution of multiple create_embedding calls."""
    embeddings_map = {
        ("model-1", ("a",)): [[1.0, 1.1, 1.2]],
        ("model-2", ("b",)): [[2.0, 2.1, 2.2]],
        ("model-3", ("c",)): [[3.0, 3.1, 3.2]],
    }
    patch_openai_async_client(embeddings_map)
    client = AsyncOpenAIClientV1(api_key="dummy-key")
    # Run three calls concurrently
    results = await asyncio.gather(
        client.create_embedding("model-1", ["a"]),
        client.create_embedding("model-2", ["b"]),
        client.create_embedding("model-3", ["c"]),
    )

@pytest.mark.asyncio
async def test_create_embedding_edge_exception(monkeypatch):
    """Test that exceptions in the async context are propagated."""
    class FailingEmbeddingsCreate:
        async def create(self, model, input):
            raise ValueError("Simulated error")

    class FailingAsyncClient:
        def __init__(self, api_key=None, base_url=None):
            self.embeddings = FailingEmbeddingsCreate()

    monkeypatch.setattr(openai, "AsyncClient", lambda api_key=None, base_url=None: FailingAsyncClient(api_key, base_url))
    client = AsyncOpenAIClientV1(api_key="dummy-key")
    with pytest.raises(ValueError, match="Simulated error"):
        await client.create_embedding("any-model", ["input"])

@pytest.mark.asyncio
async def test_create_embedding_edge_non_string_input(monkeypatch, patch_openai_async_client):
    """Test with input containing non-string elements (should be handled as per function contract)."""
    # The function expects List[str], but let's see what happens if we pass ints
    embeddings_map = {
        ("model-z", (str(123),)): [[0.9, 0.8, 0.7]]
    }
    patch_openai_async_client(embeddings_map)
    client = AsyncOpenAIClientV1(api_key="dummy-key")
    # Convert int to str as the key in embeddings_map expects str
    result = await client.create_embedding("model-z", [str(123)])

# ---- LARGE SCALE TEST CASES ----

@pytest.mark.asyncio

async def test_create_embedding_large_scale_concurrent(monkeypatch, patch_openai_async_client):
    """Test concurrent execution with many parallel calls."""
    embeddings_map = {
        ("model-a", ("x",)): [[0.1, 0.2, 0.3]],
        ("model-b", ("y",)): [[0.4, 0.5, 0.6]],
    }
    patch_openai_async_client(embeddings_map)
    client = AsyncOpenAIClientV1(api_key="dummy-key")
    # Run 50 concurrent calls (well below 1000)
    tasks = [
        client.create_embedding("model-a", ["x"]) if i % 2 == 0 else client.create_embedding("model-b", ["y"])
        for i in range(50)
    ]
    results = await asyncio.gather(*tasks)
    # Should alternate between the two embeddings
    for i, res in enumerate(results):
        if i % 2 == 0:
            pass
        else:
            pass

# ---- THROUGHPUT TEST CASES ----

@pytest.mark.asyncio


async def test_create_embedding_throughput_large_load(monkeypatch, patch_openai_async_client):
    """Throughput: Test function performance under large load (500 concurrent calls)."""
    embeddings_map = {
        ("model-t", ("input",)): [[0.1, 0.2, 0.3]],
    }
    patch_openai_async_client(embeddings_map)
    client = AsyncOpenAIClientV1(api_key="dummy-key")
    tasks = [client.create_embedding("model-t", ["input"]) for _ in range(500)]
    results = await asyncio.gather(*tasks)
    for res in results:
        pass

@pytest.mark.asyncio
async def test_create_embedding_throughput_varied_inputs(monkeypatch, patch_openai_async_client):
    """Throughput: Test with varied input sizes and concurrent calls."""
    embeddings_map = {
        ("model-t", ("a",)): [[0.1, 0.2, 0.3]],
        ("model-t", ("b",)): [[0.4, 0.5, 0.6]],
        ("model-t", ("c",)): [[0.7, 0.8, 0.9]],
    }
    patch_openai_async_client(embeddings_map)
    client = AsyncOpenAIClientV1(api_key="dummy-key")
    tasks = [
        client.create_embedding("model-t", ["a"]),
        client.create_embedding("model-t", ["b"]),
        client.create_embedding("model-t", ["c"]),
    ] * 50  # 150 concurrent calls
    results = await asyncio.gather(*tasks)
    # Should cycle through the three embeddings
    for i, res in enumerate(results):
        idx = i % 3
        expected = [[0.1, 0.2, 0.3]] if idx == 0 else ([[0.4, 0.5, 0.6]] if idx == 1 else [[0.7, 0.8, 0.9]])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.